// Copyright 2022 jeff.li. and/or its affiliates.


#include <tbir/pipeline/jit_object.h>
#include <tbir/runtime/container.h>
#include <tbir/runtime/container_private.h>
#include <tbir/runtime/ft_container.h>
#include <tbir/runtime/registry.h>

namespace tbir::runtime {

    TBIR_REGISTER_GLOBAL("runtime.JitObject_GetSelf").set_body([](PyArgs args) -> RTValue {
        MXCHECK_EQ(args.size(), 1) << "[JitObject_GetSelf] Expect 1 arguments but get " << args.size();
        UserDataRef ud_ref = args[0].As<UserDataRef>();
        auto jit_obj_ptr = check_get_jit_object(ud_ref);
        return jit_obj_ptr->self();
    });

    TBIR_REGISTER_GLOBAL("runtime.JitObject_GetFunction").set_body([](PyArgs args) -> RTValue {
        MXCHECK_EQ(args.size(), 2) << "[JitObject_GetFunction] Expect 2 arguments but get "
                                   << args.size();
        UserDataRef ud_ref = args[0].As<UserDataRef>();
        auto jit_obj_ptr = check_get_jit_object(ud_ref);
        MXCHECK(args[1].IsUnicode() || args[1].IsString());
        String name;
        if (args[1].IsUnicode()) {
            name = args[1].As<Unicode>().encode();
        } else {
            name = args[1].As<String>();
        }
        auto ret = jit_obj_ptr->GetFunction(name);
        return RTValue(new NativeFunction(ret.first), TypeIndex::kRuntimePackedFuncHandle);
    });

    TBIR_REGISTER_GLOBAL("runtime.JitObject_GetFunctionSchema")
    .set_body([](PyArgs args) -> RTValue {
        MXCHECK_EQ(args.size(), 2) << "[JitObject_GetFunction] Expect 2 arguments but get "
                                   << args.size();
        UserDataRef ud_ref = args[0].As<UserDataRef>();
        auto jit_obj_ptr = check_get_jit_object(ud_ref);
        MXCHECK(args[1].IsUnicode() || args[1].IsString());
        String name;
        if (args[1].IsUnicode()) {
            name = args[1].As<Unicode>().encode();
        } else {
            name = args[1].As<String>();
        }
        auto ret = jit_obj_ptr->GetFunction(name);
        return ret.second->ToDict();
    });

}  // namespace tbir::runtime
